import tkinter as tk
from tkinter import scrolledtext
from zhipuai import ZhipuAI

# 填写您自己的APIKey
client = ZhipuAI(api_key="")  

# 生成角色人设
def generate_personas(text):
    prompt = f"根据以下文本,为两个主要角色生成详细的角色人设:\n\n{text}"
    response = client.chat.completions.create(
        model="glm-4",
        prompt=prompt,
        max_tokens=500,
        n=1,
        stop=None,
        temperature=0.7,
    )
    return response.choices[0].text.split("\n\n")

# 生成对话
def generate_dialogue(persona1, persona2):
    char1 = ZhipuAI.create(persona1)
    char2 = ZhipuAI.create(persona2)

    dialogue = []
    for i in range(10):
        if i % 2 == 0:
            response = char1.generate_utterance()
        else:
            response = char2.generate_utterance(dialogue)
        dialogue.append(response)
    return "\n".join(dialogue)

# 保存对话到文件
def save_dialogue(dialogue, filename):
    with open(filename, "w", encoding="utf-8") as f:
        f.write(dialogue)

# 生成角色人设和对话
def generate_personas_and_dialogue(text):
    personas = generate_personas(text)
    dialogue = generate_dialogue(personas[0], personas[1])
    return "\n\n".join(personas), dialogue


# 创建主窗口
root = tk.Tk()
root.title("Role-Play Dialogue Generator")

# 创建文本输入区域
text_label = tk.Label(root, text="Input Text:")
text_label.pack()
text_area = scrolledtext.ScrolledText(root, wrap=tk.WORD, width=40, height=10)
text_area.pack()

# 创建角色人设显示区域
persona_label = tk.Label(root, text="Character Personas:")
persona_label.pack()
persona_display = scrolledtext.ScrolledText(
    root, wrap=tk.WORD, width=40, height=10)
persona_display.pack()

# 创建对话显示区域
dialogue_label = tk.Label(root, text="Dialogue:")
dialogue_label.pack()
dialogue_display = scrolledtext.ScrolledText(
    root, wrap=tk.WORD, width=40, height=10)
dialogue_display.pack()

# 生成角色人设和对话
def generate_dialogue_ui():
    text = text_area.get("1.0", tk.END).strip()
    personas, dialogue = generate_personas_and_dialogue(text)
    persona_display.delete("1.0", tk.END)
    persona_display.insert(tk.END, personas)
    dialogue_display.delete("1.0", tk.END)
    dialogue_display.insert(tk.END, dialogue)
    save_dialogue(dialogue, "dialogue.txt")


# 创建生成对话按钮
generate_button = tk.Button(
    root, text="Generate Dialogue", command=generate_dialogue_ui)
generate_button.pack()

# 运行主循环
root.mainloop()
